/*
*
客户端http审计代码
*/
package audit

import (
	"github.com/fatedier/frp/pkg/util/log"
	"net/http"
	"strconv"
	"strings"
)

const max_bs_size = 1024 * 1024 * 10

type HttpAudit struct {
	requestHeader  http.Header
	responseHeader http.Header
	responseBody   string
	audit          bool //是否审计
}

// MyResponseWriter 是一个自定义的 ResponseWriter
type HttpResponseWriter struct {
	http.ResponseWriter
	statusCode int
	bs         []byte
	bsFlag     bool
}

// HttpResponseWriter 创建一个新的 MyResponseWriter
func NewHttpResponseWriter(w http.ResponseWriter) *HttpResponseWriter {
	return &HttpResponseWriter{ResponseWriter: w, statusCode: http.StatusOK, bs: make([]byte, 0), bsFlag: true}
}

// WriteHeader 重写 WriteHeader 方法以记录状态码
func (mrw *HttpResponseWriter) WriteHeader(statusCode int) {
	mrw.statusCode = statusCode
	mrw.ResponseWriter.WriteHeader(statusCode)
}

// Write 重写 Write 方法以在写入时做一些额外的处理
func (mrw *HttpResponseWriter) Write(b []byte) (int, error) {
	// 在这里，你可以在写入响应体之前或之后执行额外的逻辑
	// 例如，记录日志、修改响应体等
	// ...

	// 调用原始的 Write 方法
	if mrw.bsFlag {
		mrw.bs = append(mrw.bs, b...)
		//log.Infof("send data.size=%d", len(b))
		if len(mrw.bs) >= max_bs_size {
			mrw.bs = make([]byte, 0)
			mrw.bsFlag = false
		}
	}
	n, err := mrw.ResponseWriter.Write(b)

	// 可以在这里添加更多的逻辑
	// ...

	return n, err
}

func (mrw *HttpResponseWriter) Analyze(r *http.Request) HttpAudit {
	HttpAudit := HttpAudit{}
	contentType := mrw.Header().Get("Content-Type")
	contentLength := mrw.Header().Get("Content-Length")
	log.Infof("frp client proxy send url=%s,Host=%s,Method=%s,status=%d contentType=%s,contentLength=%s,bs.size=%d", r.RequestURI, r.Host, r.Method, mrw.statusCode, contentType, contentLength, len(mrw.bs))
	contentLen, _ := strconv.Atoi(contentLength)
	if contentLen == len(mrw.bs) && (strings.Contains(contentType, "html") || strings.Contains(contentType, "javascript") || strings.Contains(contentType, "text") || strings.Contains(contentType, "json") || strings.Contains(contentType, "xml ")) {
		HttpAudit.responseBody = string(mrw.bs)
	}
	HttpAudit.requestHeader = r.Header
	HttpAudit.responseHeader = r.Header
	HttpAudit.audit = true
	return HttpAudit
}
